import torch

from cnn_mnist import my_model, utils

import matplotlib.pyplot as plt

from cnn_mnist.utils import plt_util
from cnn_mnist.utils.tool import predict_digit


def load_model(model_path, device='cuda'):
    le_net5 = my_model.LeNet5()

    try:
        checkpoint = torch.load(
            model_path,
            weights_only=True,
            map_location=torch.device(device)
        )

        le_net5.load_state_dict(checkpoint)
        print("模型加载成功")
    except Exception as e:
        print(f"加载模型失败: {e}")
        print("使用随机初始化参数")

    le_net5.to(device)
    le_net5.eval()

    return le_net5


def main(device='cuda'):
    model_path = "lenet5_mnist_fully_sequential.pth"

    # 检查设备
    if device == 'cuda' and not torch.cuda.is_available():
        print("CUDA不可用，将使用CPU")
        device = 'cpu'

    image_path = "./number/999.jpg"

    le_net5 = load_model(model_path, device)

    input_tensor = plt_util.preprocess_image(image_path, show_steps=True)

    input_tensor = input_tensor.to(device)

    predicted_digit, confidence = predict_digit(le_net5, input_tensor, device)

    if predicted_digit is not None:
        # 显示预测结果
        print(f"\n{'=' * 30}")
        print(f"预测结果: {predicted_digit}")
        print(f"置信度: {confidence:.4f} ({confidence * 100:.2f}%)")
        print(f"{'=' * 30}")

        # 可视化预测概率
        plt.figure(figsize=(10, 4))
        plt.bar(range(10), torch.nn.functional.softmax(le_net5(input_tensor), dim=1)[0].detach().cpu().numpy())
        plt.xticks(range(10))
        plt.xlabel('数字类别')
        plt.ylabel('预测概率')
        plt.title('各类别的预测概率分布')
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()
    else:
        print("预测失败：无法处理图像")


# 验证实现的正确性
if __name__ == "__main__":
    main()
